# koroFileHeader at Yungoal acer
# Create: 2021-06-10 11:52:48
# LastEdit: 2024-07-26 17:23:16
"""检查备份作业是否完成（Completed）
Portal 上能查到很久之前的历史数据，但是
Azure Resource Graph 浏览器 里只能查到最近2周的记录
"""
__author__ = '749B'

import os, sys
BASE_DIR = os.getcwd()
sys.path.append(BASE_DIR)

import argparse
from dateutil.parser import parse
from datetime import timezone
import textwrap
import msrestazure.azure_cloud
# pylint: disable = import-error
from azure_client import Azure_Client
from azure_client import data_format
from utils.blueking import job_start, job_success, job_fail
from utils.env import init_environs

from typing import Optional, List


def str_query_base(
    client: Azure_Client, 
    start: Optional[str], 
    end: Optional[str],
    vault: Optional[str],
    workload: Optional[str],
    ) -> str:
    """生成查询语句"""
    subscriptionDisplayName = client.subscription_display_name()
    strQuery = textwrap.dedent("""\
        recoveryservicesresources
        | where type == "microsoft.recoveryservices/vaults/backupjobs"
        {subscriptionDisplayName}
        | extend 
            vaultName = split(id, '/')[8],
            workload = properties.entityFriendlyName,
            operation = properties.operation,
            status = properties.status,
            jobType = properties.jobType,
            startTime = properties.startTime,
            duration = properties.duration
        | project subscriptionDisplayName, resourceGroup, vaultName, workload, operation, status, jobType, startTime, duration
        """).format(
            subscriptionDisplayName=subscriptionDisplayName, 
        )
    if vault:
        strQuery = strQuery + '| where vaultName =~ "%s"\n' % vault
    if workload:
        strQuery = strQuery + '| where workload =~ "%s"\n' % workload
    if start:
        start_utc = parse(start).astimezone(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
        strQuery = strQuery + '| where startTime >= datetime("%s")\n' % start_utc
    if end:
        end_utc = parse(end).astimezone(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
        strQuery = strQuery + '| where startTime < datetime("%s")\n' % end_utc
    return strQuery

def get_vm_count(client: Azure_Client, base_query: str) -> int:
    """返回记录的数量"""
    strQuery = base_query + '| summarize count()\n'
    res = client.resource_graph_query_default(
        strQuery,
        [sub.get("subscription_id") for sub in client.get_subscriptions()],
    )
    return res.data[0].get('count_')

def get_query(
    client: Azure_Client, 
    str_query: str,
    ) -> int:
    """查询：
    直接打印查询结果
    返回查询到的记录数量
    """
    # pylint: disable = no-member
    res = client.resource_graph_query_default(
        str_query, 
        [sub.get('subscription_id') for sub in client.get_subscriptions()], 
        )
    # 查询：无结果。
    if not res.data:
        return 0
    # UTC转为本地时间
    data_format.parse_time(res.data)
    # 打印结果
    data_format.output_csv(res.data, comma_begin=False)
    return len(res.data)

def init_argparse(args: Optional[List[str]] = None) -> argparse.Namespace:
    """初始化：命令行参数"""
    parser = argparse.ArgumentParser(description='检查备份作业')
    
    parser.add_argument('-e', '--envfile', help="存有用户认证信息的env文件的文件名")
    parser.add_argument('--start', help='Start time')
    parser.add_argument('--end', help='End time')
    parser.add_argument('--not_completed_only', action='store_true', help='只检查状态不是"Completed"的结果')
    parser.add_argument('--vault', help='限定保管库')
    parser.add_argument('--workload', help='限定workload(Backup item)')
    parser.add_argument('--show_query', action='store_true', help='只打印KQL查询语句，检查查询语句问题')

    args = parser.parse_args(args)
    return args

def main(args: Optional[List[str]] = None) -> None:
    """主函数"""
    # 初始化命令行参数
    if args is None:
        # 蓝鲸作业平台有BUG，重做的任务，空着的参数里都会被填充null
        args = [arg for arg in sys.argv[1:] if arg != 'null' and not arg.endswith('=null')]
    args = init_argparse(args)

    envs = {}
    if 'JENKINS_HOME' not in os.environ:
        if not args.envfile:
            return job_fail("请提供存有用户认证信息的env文件的文件名")
        # 初始化环境变量
        envs = init_environs(args.envfile)
        if not envs:
            return job_fail("env文件解密错误")
        if envs['runtime_env'] == 'blueking':
            for key in args.__dict__.keys():
                if key.startswith('dev_'):
                    v = getattr(args, key)
                    if v:
                        print(f'{key}={v}: 不支持的命令行参数，已经忽略。', file=sys.stderr)
                        setattr(args, key, None)

    # 实例化客户端
    CLOUD = getattr(msrestazure.azure_cloud, os.environ.get('CLOUD_NAME', envs.get('cloud_name')))
    client = Azure_Client(CLOUD)

    base_query = str_query_base(client, args.start, args.end, args.vault, args.workload)
    if args.show_query:
        print(base_query)  # 蓝鲸平台上，那么多内容放到下面的job_success里会返回错误。单独print出来。
        return job_success("已打印查询语句，程序结束。")

    # 检查记录数量
    query_count = get_vm_count(client, base_query)
    if not query_count:
        return job_fail("未匹配到任何虚拟机，建议检查过滤语句。")
    print("记录总数（包括所有任务状态）:", query_count)
        
    # 查询记录
    str_query = base_query + '| where status != "Completed"\n' if args.not_completed_only else base_query
    date_count = get_query(client, str_query)
    
    # 如果只需要检查异常结果，并且没有异常，就打印最早的一条记录
    if args.not_completed_only:
        if not date_count:
            str_query = base_query + '| order by todatetime(startTime) asc\n| limit 1\n'
            get_query(client, str_query)
            print('')
            return job_success("没有异常记录，只打印了时间最早的一条记录")
        else:
            print('')
            return job_success("以上是状态异常的记录: %s" % date_count)

    print('')
    return job_success("所有记录打印完毕")


if __name__ == '__main__':

    job_start()

    ###### 可在此处开始编写您的脚本逻辑代码
    ###### iJobs中执行脚本成功和失败的标准只取决于脚本最后一条执行语句的返回值
    ###### 如果返回值为0，则认为此脚本执行成功，如果非0，则认为脚本执行失败

    main()

